import sys
import os
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QLabel,
                             QPushButton, QVBoxLayout, QHBoxLayout,
                             QShortcut, QFileDialog, QMessageBox)
from PyQt5.QtCore import Qt, QPoint, QRect, QSize
from PyQt5.QtGui import (QPainter, QPen, QColor, QCursor, QBrush,
                         QPixmap, QImage, QIcon, QKeySequence)


class ScreenShotWidget(QWidget):
    """截图画布组件"""

    def __init__(self, parent=None):
        super().__init__(parent)
        # 屏幕设置
        self.setWindowFlags(Qt.FramelessWindowHint | Qt.WindowStaysOnTopHint)
        self.setWindowState(Qt.WindowFullScreen)
        self.setAttribute(Qt.WA_TranslucentBackground)

        # 截图状态
        self.is_drawing = False
        self.start_point = QPoint()
        self.end_point = QPoint()
        self.capture_rect = QRect()

        # 标注状态
        self.annotations = []
        self.current_annotation = None
        self.annotation_mode = None  # 'rectangle', 'arrow', 'text'
        self.text_edit = None

        # 画笔设置
        self.pen = QPen(QColor(255, 0, 0), 2, Qt.SolidLine)
        self.brush = QBrush(QColor(255, 255, 0, 100))

        # 获取屏幕截图
        screen = QApplication.primaryScreen()
        self.screen_pixmap = screen.grabWindow(0)

        # 光标设置
        self.setCursor(Qt.CrossCursor)

    def paintEvent(self, event):
        """绘制事件"""
        painter = QPainter(self)
        # 绘制屏幕截图
        painter.drawPixmap(0, 0, self.screen_pixmap)

        # 绘制灰色遮罩层
        painter.setPen(Qt.NoPen)
        painter.setBrush(QColor(0, 0, 0, 100))
        painter.drawRect(0, 0, self.width(), self.height())

        # 绘制选区（无遮罩）
        if not self.capture_rect.isNull():
            painter.setCompositionMode(QPainter.CompositionMode_Source)
            painter.drawPixmap(self.capture_rect, self.screen_pixmap, self.capture_rect)
            painter.setCompositionMode(QPainter.CompositionMode_SourceOver)

            # 绘制选区边框
            painter.setPen(QPen(QColor(0, 150, 255), 2, Qt.DashLine))
            painter.drawRect(self.capture_rect)

            # 显示选区尺寸
            painter.setPen(QPen(Qt.white, 1))
            painter.setFont(self.font())
            rect_text = f"{self.capture_rect.width()} x {self.capture_rect.height()}"
            painter.drawText(
                self.capture_rect.bottomRight() + QPoint(-100, 20),
                rect_text
            )

        # 绘制当前标注
        if self.current_annotation:
            self.draw_annotation(painter, self.current_annotation)

        # 绘制所有已完成的标注
        for annotation in self.annotations:
            self.draw_annotation(painter, annotation)

    def draw_annotation(self, painter, annotation):
        """绘制标注"""
        if annotation['type'] == 'rectangle':
            painter.setPen(QPen(annotation['color'], annotation['width']))
            painter.setBrush(Qt.NoBrush)
            painter.drawRect(annotation['rect'])

        elif annotation['type'] == 'arrow':
            painter.setPen(QPen(annotation['color'], annotation['width']))
            start, end = annotation['points']
            painter.drawLine(start, end)

            # 绘制箭头头部
            angle = start.angleTo(end) / 16.0
            arrow_size = 10
            arrow_p1 = end - QPoint(
                arrow_size * (0.8 * angle.cos() + 0.2 * (angle + 90).sin()),
                arrow_size * (0.8 * angle.sin() - 0.2 * (angle + 90).cos())
            )
            arrow_p2 = end - QPoint(
                arrow_size * (0.8 * angle.cos() - 0.2 * (angle + 90).sin()),
                arrow_size * (0.8 * angle.sin() + 0.2 * (angle + 90).cos())
            )
            painter.drawLine(end, arrow_p1)
            painter.drawLine(end, arrow_p2)

        elif annotation['type'] == 'text':
            painter.setPen(QPen(annotation['color'], 1))
            painter.setFont(annotation['font'])
            painter.drawText(annotation['rect'], Qt.TextWordWrap, annotation['text'])

    def mousePressEvent(self, event):
        """鼠标按下事件"""
        if event.button() == Qt.LeftButton:
            if not self.annotation_mode:
                # 开始截图选区
                self.is_drawing = True
                self.start_point = event.pos()
                self.end_point = event.pos()
                self.capture_rect = QRect()
            else:
                # 开始绘制标注
                self.current_annotation = {
                    'type': self.annotation_mode,
                    'color': self.pen.color(),
                    'width': self.pen.width(),
                    'start': event.pos()
                }
                if self.annotation_mode == 'rectangle':
                    self.current_annotation['rect'] = QRect(event.pos(), event.pos())
                elif self.annotation_mode == 'arrow':
                    self.current_annotation['points'] = (event.pos(), event.pos())
                elif self.annotation_mode == 'text':
                    self.current_annotation['rect'] = QRect(event.pos(), QSize(200, 100))
                    self.current_annotation['text'] = ""
                    self.current_annotation['font'] = self.font()
                    self.show_text_editor(event.pos())

        elif event.button() == Qt.RightButton:
            # 右键取消当前操作
            if self.annotation_mode:
                self.annotation_mode = None
                self.current_annotation = None
                if self.text_edit:
                    self.text_edit.hide()
                    self.text_edit = None
            else:
                self.close()

    def mouseMoveEvent(self, event):
        """鼠标移动事件"""
        if self.is_drawing:
            # 更新截图选区
            self.end_point = event.pos()
            self.capture_rect = QRect(
                min(self.start_point.x(), self.end_point.x()),
                min(self.start_point.y(), self.end_point.y()),
                abs(self.start_point.x() - self.end_point.x()),
                abs(self.start_point.y() - self.end_point.y())
            )
            self.update()

        elif self.current_annotation:
            # 更新当前标注
            if self.annotation_mode == 'rectangle':
                self.current_annotation['rect'] = QRect(
                    self.current_annotation['start'],
                    event.pos()
                ).normalized()
            elif self.annotation_mode == 'arrow':
                self.current_annotation['points'] = (
                    self.current_annotation['start'],
                    event.pos()
                )
            self.update()

    def mouseReleaseEvent(self, event):
        """鼠标释放事件"""
        if event.button() == Qt.LeftButton:
            if self.is_drawing:
                # 完成截图选区
                self.is_drawing = False
                self.end_point = event.pos()
                self.capture_rect = QRect(
                    min(self.start_point.x(), self.end_point.x()),
                    min(self.start_point.y(), self.end_point.y()),
                    abs(self.start_point.x() - self.end_point.x()),
                    abs(self.start_point.y() - self.end_point.y())
                )
                self.update()

            elif self.current_annotation and self.annotation_mode != 'text':
                # 完成标注（文本标注在编辑器完成）
                self.annotations.append(self.current_annotation)
                self.current_annotation = None
                self.annotation_mode = None
                self.update()

    def show_text_editor(self, pos):
        """显示文本编辑器"""
        if self.text_edit:
            self.text_edit.hide()

        self.text_edit = QLabel(self)
        self.text_edit.setGeometry(pos.x(), pos.y(), 200, 100)
        self.text_edit.setStyleSheet("""
            background-color: rgba(255, 255, 200, 220);
            border: 1px solid #FF9900;
            padding: 5px;
        """)
        self.text_edit.setText("输入文本...")
        self.text_edit.setAlignment(Qt.AlignTop | Qt.AlignLeft)
        self.text_edit.setWordWrap(True)
        self.text_edit.mousePressEvent = self.finish_text_annotation
        self.text_edit.show()

    def finish_text_annotation(self, event):
        """完成文本标注"""
        if self.current_annotation and self.text_edit:
            self.current_annotation['text'] = self.text_edit.text()
            self.current_annotation['font'] = self.text_edit.font()
            self.annotations.append(self.current_annotation)
            self.current_annotation = None
            self.annotation_mode = None
            self.text_edit.hide()
            self.text_edit = None
            self.update()

    def get_cropped_image(self):
        """获取裁剪后的图像（带标注）"""
        # 创建透明背景图像
        image = QImage(self.capture_rect.size(), QImage.Format_ARGB32)
        image.fill(Qt.transparent)

        # 绘制截图和标注
        painter = QPainter(image)
        # 绘制原始截图
        painter.drawPixmap(0, 0, self.screen_pixmap.copy(self.capture_rect))

        # 绘制标注（仅限选区内的标注）
        for annotation in self.annotations:
            # 将标注坐标转换为相对于选区的坐标
            if annotation['type'] == 'rectangle':
                rect = annotation['rect'].intersected(self.capture_rect)
                if not rect.isNull():
                    adj_rect = rect.translated(-self.capture_rect.topLeft())
                    painter.setPen(QPen(annotation['color'], annotation['width']))
                    painter.setBrush(Qt.NoBrush)
                    painter.drawRect(adj_rect)

            elif annotation['type'] == 'arrow':
                start, end = annotation['points']
                if self.capture_rect.contains(start) and self.capture_rect.contains(end):
                    adj_start = start - self.capture_rect.topLeft()
                    adj_end = end - self.capture_rect.topLeft()
                    painter.setPen(QPen(annotation['color'], annotation['width']))
                    painter.drawLine(adj_start, adj_end)

            elif annotation['type'] == 'text':
                if self.capture_rect.intersects(annotation['rect']):
                    adj_rect = annotation['rect'].translated(-self.capture_rect.topLeft())
                    painter.setPen(QPen(annotation['color'], 1))
                    painter.setFont(annotation['font'])
                    painter.drawText(adj_rect, Qt.TextWordWrap, annotation['text'])

        painter.end()
        return image


class ScreenShotTool(QMainWindow):
    """主截图工具窗口"""

    def __init__(self):
        super().__init__()
        self.setWindowTitle("高级截图工具")
        self.setGeometry(100, 100, 400, 200)
        self.setWindowIcon(QIcon(self.create_icon()))

        # 创建主组件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)

        # 创建布局
        layout = QVBoxLayout(central_widget)
        layout.setAlignment(Qt.AlignCenter)

        # 标题
        title = QLabel("PyQt5 高级截图工具")
        title.setStyleSheet("font-size: 18px; font-weight: bold; margin-bottom: 20px;")
        title.setAlignment(Qt.AlignCenter)
        layout.addWidget(title)

        # 按钮区域
        btn_layout = QHBoxLayout()

        # 截图按钮
        self.capture_btn = QPushButton("开始截图")
        self.capture_btn.setStyleSheet("""
            QPushButton {
                background-color: #4CAF50;
                color: white;
                font-weight: bold;
                border-radius: 5px;
                padding: 10px 20px;
                font-size: 14px;
            }
            QPushButton:hover {
                background-color: #45a049;
            }
        """)
        self.capture_btn.setCursor(Qt.PointingHandCursor)
        self.capture_btn.clicked.connect(self.start_screenshot)
        btn_layout.addWidget(self.capture_btn)

        # 退出按钮
        exit_btn = QPushButton("退出")
        exit_btn.setStyleSheet("""
            QPushButton {
                background-color: #f44336;
                color: white;
                font-weight: bold;
                border-radius: 5px;
                padding: 10px 20px;
                font-size: 14px;
            }
            QPushButton:hover {
                background-color: #d32f2f;
            }
        """)
        exit_btn.setCursor(Qt.PointingHandCursor)
        exit_btn.clicked.connect(self.close)
        btn_layout.addWidget(exit_btn)

        layout.addLayout(btn_layout)

        # 添加快捷键
        QShortcut(QKeySequence("Ctrl+Shift+S"), self, self.start_screenshot)
        QShortcut(QKeySequence("Ctrl+Q"), self, self.close)

    def create_icon(self):
        """创建应用图标"""
        pixmap = QPixmap(64, 64)
        pixmap.fill(Qt.transparent)

        painter = QPainter(pixmap)
        painter.setBrush(QColor(66, 133, 244))
        painter.drawRect(10, 10, 44, 44)

        painter.setPen(QPen(Qt.white, 3))
        painter.drawLine(20, 20, 44, 44)
        painter.drawLine(44, 20, 20, 44)

        painter.end()
        return pixmap

    def start_screenshot(self):
        """开始截图"""
        self.hide()
        self.screenshot_widget = ScreenShotWidget()
        self.screenshot_widget.show()

        # 连接完成信号
        # self.screenshot_widget.keyPressEvent = self.handle_screenshot_keys

    def handle_screenshot_keys(self, event):
        """处理截图窗口的按键事件"""
        if event.key() == Qt.Key_Escape:
            # ESC 取消截图
            self.screenshot_widget.close()
            self.show()

        elif event.key() == Qt.Key_Enter or event.key() == Qt.Key_Return:
            # Enter 完成截图
            self.complete_screenshot()

        elif event.key() == Qt.Key_R:
            # R 切换到矩形标注模式
            self.screenshot_widget.annotation_mode = 'rectangle'
            self.screenshot_widget.setCursor(Qt.CrossCursor)

        elif event.key() == Qt.Key_A:
            # A 切换到箭头标注模式
            self.screenshot_widget.annotation_mode = 'arrow'
            self.screenshot_widget.setCursor(Qt.CrossCursor)

        elif event.key() == Qt.Key_T:
            # T 切换到文本标注模式
            self.screenshot_widget.annotation_mode = 'text'
            self.screenshot_widget.setCursor(Qt.IBeamCursor)

        elif event.key() == Qt.Key_Z and event.modifiers() == Qt.ControlModifier:
            # Ctrl+Z 撤销上一个标注
            if self.screenshot_widget.annotations:
                self.screenshot_widget.annotations.pop()
                self.screenshot_widget.update()

        elif event.key() == Qt.Key_C and event.modifiers() == Qt.ControlModifier:
            # Ctrl+C 复制到剪贴板
            self.copy_to_clipboard()

        elif event.key() == Qt.Key_S and event.modifiers() == Qt.ControlModifier:
            # Ctrl+S 保存截图
            self.save_screenshot()

    def complete_screenshot(self):
        """完成截图并显示结果"""
        if not self.screenshot_widget.capture_rect.isNull():
            # 获取截图
            image = self.screenshot_widget.get_cropped_image()

            # 关闭截图窗口
            self.screenshot_widget.close()
            self.show()

            # 显示结果窗口
            self.show_result(image)

    def show_result(self, image):
        """显示截图结果"""
        result_dialog = QWidget()
        result_dialog.setWindowTitle("截图结果")
        result_dialog.setGeometry(100, 100, 800, 600)
        result_dialog.setWindowFlags(Qt.WindowStaysOnTopHint)

        layout = QVBoxLayout(result_dialog)

        # 显示图像
        image_label = QLabel()
        image_label.setPixmap(QPixmap.fromImage(image))
        image_label.setAlignment(Qt.AlignCenter)
        layout.addWidget(image_label)

        # 按钮区域
        btn_layout = QHBoxLayout()

        # 复制按钮
        copy_btn = QPushButton("复制到剪贴板")
        copy_btn.clicked.connect(lambda: self.copy_to_clipboard(image))
        btn_layout.addWidget(copy_btn)

        # 保存按钮
        save_btn = QPushButton("保存图片")
        save_btn.clicked.connect(lambda: self.save_screenshot(image))
        btn_layout.addWidget(save_btn)

        # 关闭按钮
        close_btn = QPushButton("关闭")
        close_btn.clicked.connect(result_dialog.close)
        btn_layout.addWidget(close_btn)

        layout.addLayout(btn_layout)
        result_dialog.show()

    def copy_to_clipboard(self, image=None):
        """复制截图到剪贴板"""
        if image is None and hasattr(self, 'screenshot_widget'):
            image = self.screenshot_widget.get_cropped_image()

        if image and not image.isNull():
            clipboard = QApplication.clipboard()
            clipboard.setImage(image)
            QMessageBox.information(
                self, "成功",
                "截图已复制到剪贴板！\n您可以直接粘贴到其他应用中。"
            )

    def save_screenshot(self, image=None):
        """保存截图到文件"""
        if image is None and hasattr(self, 'screenshot_widget'):
            image = self.screenshot_widget.get_cropped_image()

        if image and not image.isNull():
            file_path, _ = QFileDialog.getSaveFileName(
                self, "保存截图",
                os.path.expanduser("~/screenshot.png"),
                "PNG 图片 (*.png);;JPEG 图片 (*.jpg *.jpeg)"
            )

            if file_path:
                if file_path.lower().endswith(('.png', '.jpg', '.jpeg')):
                    image.save(file_path)
                    QMessageBox.information(self, "成功", f"截图已保存到:\n{file_path}")
                else:
                    # 如果没有扩展名，添加.png
                    file_path += ".png"
                    image.save(file_path)
                    QMessageBox.information(self, "成功", f"截图已保存到:\n{file_path}")


if __name__ == "__main__":
    app = QApplication(sys.argv)
    app.setStyle("Fusion")  # 使用Fusion风格使界面更美观

    # 设置全局字体
    font = app.font()
    font.setPointSize(10)
    app.setFont(font)

    window = ScreenShotTool()
    window.show()
    sys.exit(app.exec_())